from auton_survival.estimators import SurvivalModel
from auton_survival.metrics import survival_regression_metric
from sklearn.model_selection import ParameterGrid
from estimators_demo_utils import plot_performance_metrics
import pandas as pd 
from collections import defaultdict
from auton_survival import DeepCoxPH,DeepSurvivalMachines, DeepCoxMixtures,DeepRecurrentCoxPH
from sksurv.metrics import concordance_index_ipcw, brier_score, cumulative_dynamic_auc
from auton_survival.estimators import SurvivalModel
from auton_survival.metrics import survival_regression_metric
from sklearn.model_selection import ParameterGrid
from collections import defaultdict
from data import give_target
from train_model import run_cox
import pandas as pd 
import numpy as np 
import os 

# 读取特征数据
panel = pd.read_csv('~/PANEL.csv')
features = panel.columns.drop(['eid','Region_code'])

# 心血管疾病大类
disease = pd.read_csv('~/all_enamd_ref.csv',low_memory=False)

disease_dict = {'DM':['E10','E11','E12','E13','E14']}
disease_icd = disease_dict['DM']   # dibetes
target_data = give_target(disease,disease_icd)  # 得到event和duration列
mydf = pd.merge(panel, target_data, on='eid')
mydf.rename(columns={'duration':'time'},inplace=True)
mydf['DM_ETH'] = mydf['DM_ETH_W'] + mydf['DM_ETH_A']*2 + mydf['DM_ETH_B']*3 + mydf['DM_ETH_O']*4 - 1
mydf = mydf.sample(n=10000, random_state=42)
# 选取特征
with open(file='~/PANEL copy.txt') as f:
    features = f.read()
features = features.split(' + ')

# 记录全部模型结果
model_results = defaultdict(list)
for model_name in ['dcph', 'dsm', 'dcm','rsf']:
    results = defaultdict(list)
    # 初始化模型
    model = SurvivalModel(model=model_name,random_seed=2024)
    for fold_id in range(10):
        train_idx = mydf['Region_code'].index[mydf['Region_code'] != fold_id]
        test_idx = mydf['Region_code'].index[mydf['Region_code'] == fold_id]
        mydf_train = mydf.loc[train_idx]
        mydf_test = mydf.loc[test_idx]
        x_tr = mydf_train[features]
        x_te = mydf_test[features]
        y_tr = mydf_train[['event','time']]
        y_te = mydf_test[['event','time']]
        #times = [5, 10, y_tr[y_tr['event']==1]['time'].max()]
        # 确定 times 列表
        # 获取 y_te 和 y_tr 的最大时间
        max_time_te = y_te[y_te['event']==1]['time'].max()
        max_time_tr = y_tr[y_tr['event']==1]['time'].max()

        # 确保 times 中的最大时间点在 y_te 和 y_tr 的最大时间范围内
        max_time = min(max_time_te, max_time_tr)
        times = [5, 10, max_time]
        # Fit the model 
        model.fit(x_tr, y_tr)
        # Obtain survival probabilities for test set
        predictions_te = model.predict_survival(x_te, times)

        # Compute the Brier Score and time-dependent concordance index for the test set to assess model performance
        c_index_result = survival_regression_metric('ctd', outcomes=y_te, predictions=predictions_te, 
                                                            times=times, outcomes_train=y_tr)
        for i, c_index in enumerate(c_index_result):
            results[i].append(c_index)
    mean_c_index_5 = np.mean(results[0])
    std_c_index_5 = np.std(results[0])
    mean_c_index_10 = np.mean(results[1])
    std_c_index_10 = np.std(results[1])
    mean_c_index_now = np.mean(results[2])
    std_c_index_now = np.std(results[2])
    interval_c_index_5 = [(mean_c_index_5 - std_c_index_5).round(2), (mean_c_index_5 + std_c_index_5).round(2)]
    interval_c_index_10 = [(mean_c_index_10 - std_c_index_10).round(2), (mean_c_index_10 + std_c_index_10).round(2)]
    interval_c_index_now = [(mean_c_index_now - std_c_index_now).round(2), (mean_c_index_now + std_c_index_now).round(2)]
    model_results['disease'].append('DM')
    model_results['model'].append(model_name)
    model_results['c_index_5'].append(mean_c_index_5)
    model_results['c_index_10'].append(mean_c_index_10)
    model_results['c_index_now'].append(mean_c_index_now)
    model_results['interval_5'].append(interval_c_index_5)
    model_results['interval_10'].append(interval_c_index_10)
    model_results['interval_now'].append(interval_c_index_now)
model_df = pd.DataFrame(model_results)
model_df.to_csv('~/C_index_enamd_all_model_1w.csv',index=False)